#include <hash/topk.h>
#include <hash/dataset.h>
#include <hash/LSH.h>
#include <hash/composite.h>
#include <hash/LSH_index.h>
#include <hash/hash_table.h>
#include <hash/eval.h>
#include <hash/bruteforce.h>
#include <hash/multiprobeLSH.h>
#include <hash/LSHforest.h>
#include <limits>
#include <stdio.h>
#include <stdlib.h>
#include <time.h>

using namespace cspace_learning;

void lsh_knn(const Dataset& dataset, const Dataset& queries);
void mpllsh_knn(const Dataset& dataset, const Dataset& queries);
void lshforest_knn(const Dataset& dataset, const Dataset& queries);
void bruteforce_knn(const Dataset& dataset, const Dataset& queries);
void generate_data(Dataset& dataset, unsigned int N, unsigned int dim);

int main()
{
  srand(time(NULL));

  // construct dataset
  // dimension 10, number 100000;
  // query 100
  unsigned int dim = 10;
  unsigned int N = 100000;
  unsigned int Q = 100;
  Dataset dataset(dim);
  Dataset queries(dim);
  generate_data(dataset, N, dim);
  generate_data(queries, Q, dim);

  bruteforce_knn(dataset, queries);
  std::cout << "==========================================" << std::endl;
  lsh_knn(dataset, queries);
  std::cout << "==========================================" << std::endl;
  mpllsh_knn(dataset, queries);
  std::cout << "==========================================" << std::endl;
  //lshforest_knn(dataset, queries);

  return 1;
}

void generate_data(Dataset& dataset, unsigned int N, unsigned int dim)
{
  float ub = 10, lb = 0;

  float* v = new float[dim];

  for(unsigned int i = 0; i < N; ++i)
  {
    for(unsigned int d = 0; d < dim; ++d)
      v[d] = (ub - lb) * (rand() / (float)RAND_MAX) + lb;
    dataset.insert(v, dim);
  }

  delete [] v; v = NULL;
}

void mpllsh_knn(const Dataset& dataset, const Dataset& queries)
{
  typedef MultiProbeLSH MyLSH;
  typedef LinearHashTable<unsigned int, unsigned int, __gnu_cxx::hash<unsigned int> > HashTable;
  typedef MultiProbeLSHIndex<HashTable, unsigned int> Index;

  Dataset::Accessor accessor(dataset);
  Index index;

  unsigned int M = 5;
  unsigned int L = 200;
  unsigned int K = 50;
  unsigned int H = 1000;
  float W = 2.0;
  float R = std::numeric_limits<float>::max();
  DefaultRng rng;
  metric::l2sqr<float> l2(dataset.dim());
  eval::PointEval<metric::l2sqr<float> > eval(l2);
  unsigned int T = 5; // multiprobe number

  // compute lsh index
  Index::Parameter param;
  param.repeat = M;
  param.range = H;
  param.W = W;
  param.dim = dataset.dim();

  index.init(param, rng, L);

  for(unsigned int i = 0; i < dataset.size(); ++i)
  {
    // insert an item to the hash table
    index.insert(i, dataset[i]);
  }

  TopkScanner<Dataset::Accessor, const float*, eval::PointEval<metric::l2sqr<float> > > scanner(accessor, &eval, K, R);

  for(unsigned int i = 0; i < queries.size(); ++i)
  {
    accessor.reset();
    eval.reset(queries[i]);
    scanner.reset(queries[i]);
    index.query(T, scanner);
    for(unsigned int j = 0; j < scanner.topk().size(); ++j)
    {
      std::cout << sqrt(scanner.topk()[j].dist) << " ";
    }
    std::cout << std::endl;
  }
}

void lsh_knn(const Dataset& dataset, const Dataset& queries)
{
  typedef Tail<RepeatHash<GaussianLSH> > MyLSH;
  typedef LinearHashTable<unsigned int, unsigned int, __gnu_cxx::hash<unsigned int> > HashTable;
  typedef LSHIndex<MyLSH, HashTable, unsigned int> Index;

  Dataset::Accessor accessor(dataset);
  Index index;

  unsigned int M = 5;
  unsigned int L = 200;
  unsigned int K = 50;
  unsigned int H = 1000;
  float W = 2.0;
  float R = std::numeric_limits<float>::max();
  DefaultRng rng;
  metric::l2sqr<float> l2(dataset.dim());
  eval::PointEval<metric::l2sqr<float> > eval(l2);

  // compute lsh index
  Index::Parameter param;
  param.repeat = M;
  param.range = H;
  param.W = W;
  param.dim = dataset.dim();

  index.init(param, rng, L);

  for(unsigned int i = 0; i < dataset.size(); ++i)
  {
    // insert an item to the hash table
    index.insert(i, dataset[i]);
  }

  TopkScanner<Dataset::Accessor, const float*, eval::PointEval<metric::l2sqr<float> > > scanner(accessor, &eval, K, R);

  for(unsigned int i = 0; i < queries.size(); ++i)
  {
    accessor.reset();
    eval.reset(queries[i]);
    scanner.reset(queries[i]);
    index.query(scanner);
    for(unsigned int j = 0; j < scanner.topk().size(); ++j)
    {
      std::cout << sqrt(scanner.topk()[j].dist) << " ";
    }
    std::cout << std::endl;
  }
}

void lshforest_knn(const Dataset& dataset, const Dataset& queries)
{
  typedef Tail<RepeatHash<GaussianLSH> > MyLSH;
  typedef ForestIndex<MyLSH, unsigned int> Index;

  Dataset::Accessor accessor(dataset);
  Index index;

  unsigned int M = 5;
  unsigned int L = 50;
  unsigned int K = 50;
  unsigned int H = 1000;
  float W = 2.0;
  float R = std::numeric_limits<float>::max();
  DefaultRng rng;
  metric::l2sqr<float> l2(dataset.dim());
  eval::PointEval<metric::l2sqr<float> > eval(l2);

  // compute lsh index
  Index::Parameter param;
  param.repeat = M;
  param.range = H;
  param.W = W;
  param.dim = dataset.dim();

  index.init(param, rng, L, M);

  for(unsigned int i = 0; i < dataset.size(); ++i)
  {
    // insert an item to the index
    index.insert(i, dataset[i]);
  }

  TopkScanner<Dataset::Accessor,  const float*, eval::PointEval<metric::l2sqr<float> > > scanner(accessor, &eval, K, R);

  unsigned int M_forest = 1000;
  for(unsigned int i = 0; i < queries.size(); ++i)
  {
    accessor.reset();
    eval.reset(queries[i]);
    scanner.reset(queries[i]);
    index.query(M_forest, scanner);
    for(unsigned int j = 0; j < scanner.topk().size(); ++j)
    {
      std::cout << sqrt(scanner.topk()[j].dist) << " ";
    }
    std::cout << std::endl;
  }
}

void bruteforce_knn(const Dataset& dataset, const Dataset& queries)
{
  Dataset::Accessor accessor(dataset);

  unsigned int K = 50;
  float R = std::numeric_limits<float>::max();
  metric::l2sqr<float> l2(dataset.dim());
  eval::PointEval<metric::l2sqr<float> > eval(l2);

  Bruteforce<Dataset::Accessor, eval::PointEval<metric::l2sqr<float> > > bruteforce(accessor, &eval, K, R);

  for(unsigned int i = 0; i < queries.size(); ++i)
  {
    accessor.reset();
    eval.reset(queries[i]);
    bruteforce.reset();
    bruteforce.knn();

    for(unsigned int j = 0; j < bruteforce.topk().size(); ++j)
    {
      std::cout << sqrt(bruteforce.topk()[j].dist) << " ";
    }
    std::cout << std::endl;
  }
}



